# core/nlp_model.py
from transformers import pipeline
import inventor_api as inv


class ModelCommander:
    def __init__(self):
        self.nlp = pipeline("text2text-generation", model="deepseek-ai/deepseek-coder-33b-instruct")
        self.api = inv.InventorAPI()

    def execute_command(self, text: str):
        """解析自然语言命令"""
        prompt = f"""将用户指令转换为Inventor API调用：
        示例输入："将壁厚增加到5mm"
        示例输出：{"thicken_wall(5)"}

        当前模型特征：{self.api.get_current_features()}
        用户指令：{text}"""

        code = self.nlp(prompt, max_length=100)[0]['generated_text']
        self._safe_execute(code)

    def _safe_execute(self, code: str):
        """安全执行生成的代码"""
        allowed_commands = {
            'extrude', 'revolve', 'thicken_wall', 'create_sketch'
        }

        # 代码白名单验证
        for cmd in allowed_commands:
            if cmd in code:
                try:
                    exec(f"self.api.{code}")
                    return
                except Exception as e:
                    raise RuntimeError(f"API执行失败: {str(e)}")

        raise ValueError("包含未授权操作指令")


# 使用示例
commander = ModelCommander()
commander.execute_command("在底面添加直径20mm的圆孔")